from multiprocessing import shared_memory, Pool, cpu_count
import numpy as np
import pandas as pd
from scipy.spatial import cKDTree
from tqdm import tqdm

# 辅助函数：评估每个组合
def evaluate_combination(params):
    grid_nodes = params['grid_nodes']
    fault_nodes = params['fault_nodes']
    average_distance = calculate_fault_distance(grid_nodes, fault_nodes)
    return (params['x_num'], params['y_num'], params['z_num'], average_distance)


# 辅助函数：包装器用于多进程生成网格节点
def generate_grid_nodes_wrapper(combination_points):
    combination, points, chunk_size = combination_points
    x_num, y_num, z_num = combination
    
    # 分块生成网格节点，并将所有块组合成一个数组
    combined_nodes = []
    for chunk_nodes in generate_chunks(points, x_num, y_num, z_num, chunk_size):
        combined_nodes.append(chunk_nodes)
    
    return combination, np.vstack(combined_nodes)

def generate_chunks(points, x_num, y_num, z_num, chunk_size=1000):
    # 计算边界值
    x_min, y_min, z_min = points[:, 0].min(), points[:, 1].min(), points[:, 2].min()
    x_max, y_max, z_max = points[:, 0].max(), points[:, 1].max(), points[:, 2].max()

    # 计算步长
    x_step = (x_max - x_min) / (x_num - 1) if x_num > 1 else 0
    y_step = (y_max - y_min) / (y_num - 1) if y_num > 1 else 0
    z_step = (z_max - z_min) / (z_num - 1) if z_num > 1 else 0

    total_indices = x_num * y_num * z_num
    for start in range(0, total_indices, chunk_size):
        end = min(start + chunk_size, total_indices)
        
        # 创建三维索引数组
        i, j, k = np.unravel_index(np.arange(start, end), (x_num, y_num, z_num))
        
        # 使用广播机制生成块中的节点
        chunk_nodes = np.array([
            x_min + i * x_step,
            y_min + j * y_step,
            z_min + k * z_step
        ]).T
        
        yield chunk_nodes




# 辅助函数：从Excel文件加载点数据
def load_points_from_excel(file_path):
    df = pd.read_excel(file_path, header=0)
    if df.shape[1] < 3:
        raise ValueError("Excel文件中必须包含至少三列数据，分别表示X, Y, Z坐标")

    def is_integer(num):
        return num == int(num)

    geo_points = []
    fault_points = []

    for _, row in df.iterrows():
        x, y, z = row[0], row[1], row[2]
        if is_integer(x) and is_integer(y):
            geo_points.append([x, y, z])
        else:
            fault_points.append([x, y, z])
    print(len(geo_points))
    print(len(fault_points))
    return np.array(geo_points), np.array(fault_points)

# 辅助函数：计算生成的网格节点与真实断层面节点之间的距离
def calculate_fault_distance(grid_nodes, fault_nodes):
    tree = cKDTree(fault_nodes)
    distances, _ = tree.query(grid_nodes)
    average_distance = np.mean(distances)
    return average_distance

#广播机制
def generate_grid_nodes(points, x_num, y_num, z_num):
    x_min, y_min, z_min = points[:, 0].min(), points[:, 1].min(), points[:, 2].min()
    x_max, y_max, z_max = points[:, 0].max(), points[:, 1].max(), points[:, 2].max()

    x_step = (x_max - x_min) / (x_num - 1)
    y_step = (y_max - y_min) / (y_num - 1)
    z_step = (z_max - z_min) / (z_num - 1)
    #print("meshgrid...")
    # 使用 NumPy 的 meshgrid 和 reshape 来创建所有的坐标点
    i, j, k = np.indices((x_num, y_num, z_num))
    grid_nodes = np.array([
        x_min + i * x_step,
        y_min + j * y_step,
        z_min + k * z_step
    ]).reshape(3, -1).T
    #print(f"grid_nodes_length:{len(grid_nodes)}")
    return grid_nodes

# 主要函数：调整网格生成参数，考虑三维空间，并使用多进程加速
def optimize_grid(points, fault_nodes, target_cell_size=(10, 11)):
    x_min, y_min, z_min = points[:, 0].min(), points[:, 1].min(), points[:, 2].min()
    x_max, y_max, z_max = points[:, 0].max(), points[:, 1].max(), points[:, 2].max()
    width, height, depth = x_max - x_min, y_max - y_min, z_max - z_min

    # 计算网格数量的上下限
    x_num_lower = int(np.ceil(width / target_cell_size[1]))
    x_num_upper = int(np.floor(width / target_cell_size[0]))
    y_num_lower = int(np.ceil(height / target_cell_size[1]))
    y_num_upper = int(np.floor(height / target_cell_size[0]))
    z_num_lower = int(np.ceil(depth / target_cell_size[1]))
    z_num_upper = int(np.floor(depth / target_cell_size[0]))

    if x_num_lower > x_num_upper or y_num_lower > y_num_upper or z_num_lower > z_num_upper:
        raise ValueError("给定的目标单元尺寸无法满足网格划分要求，请检查输入数据或调整目标单元尺寸")

    x_num_range = range(max(x_num_lower, 2), x_num_upper + 1)  # 至少需要2个单元
    y_num_range = range(max(y_num_lower, 2), y_num_upper + 1)  # 至少需要2个单元
    z_num_range = range(max(z_num_lower, 2), z_num_upper + 1)  # 至少需要2个单元
    print(f"x_num_range:{x_num_range},y_num_range:{y_num_range},z_num_range:{z_num_range}")
    combinations = [(x_num, y_num, z_num) for x_num in x_num_range for y_num in y_num_range for z_num in z_num_range]
    with Pool(cpu_count()) as pool:
        combinations_with_points = [(comb, points, 1000) for comb in combinations]  # 添加分块大小
        precomputed_grid_nodes = {}
        
        for comb, nodes in tqdm(pool.imap_unordered(generate_grid_nodes_wrapper, combinations_with_points), total=len(combinations), desc="Precomputing Grid Nodes"):
            precomputed_grid_nodes[comb] = nodes

        # 构建参数列表用于评估每个组合
        params_list = [{'grid_nodes': precomputed_grid_nodes[comb], 'fault_nodes': fault_nodes, 'x_num': comb[0], 'y_num': comb[1], 'z_num': comb[2]} for comb in combinations]

        # 使用相同的 Pool 实例来评估每个组合
        results = list(tqdm(pool.imap(evaluate_combination, params_list), total=len(combinations), desc="Evaluating Combinations"))

    best_x_num, best_y_num, best_z_num, min_average_distance = min(results, key=lambda r: r[3])

    return best_x_num, best_y_num, best_z_num, min_average_distance
# 构建六面体单元，考虑三维空间
def construct_hexahedral_grid(x_num, y_num, z_num):
    cells = []

    for i in range(x_num - 1):
        for j in range(y_num - 1):
            for k in range(z_num - 1):
                p1 = i * y_num * z_num + j * z_num + k
                p2 = (i + 1) * y_num * z_num + j * z_num + k
                p3 = (i + 1) * y_num * z_num + (j + 1) * z_num + k
                p4 = i * y_num * z_num + (j + 1) * z_num + k
                p5 = p1 + (x_num - 1) * y_num * z_num // z_num
                p6 = p2 + (x_num - 1) * y_num * z_num // z_num
                p7 = p3 + (x_num - 1) * y_num * z_num // z_num
                p8 = p4 + (x_num - 1) * y_num * z_num // z_num
                
                cells.append([p1, p2, p3, p4, p5, p6, p7, p8])

    return np.array(cells)

# 导出为 VTK 文件
def export_to_vtk(points, cells, output_file):
    with open(output_file, "w") as f:
        f.write("# vtk DataFile Version 3.0\n")
        f.write("Hexahedral grid\n")
        f.write("ASCII\n")
        f.write("DATASET UNSTRUCTURED_GRID\n")
        
        f.write(f"POINTS {len(points)} float\n")
        for point in points:
            f.write(f"{point[0]} {point[1]} {point[2]}\n")
        
        f.write(f"CELLS {len(cells)} {len(cells) * 9}\n")
        for cell in cells:
            f.write(f"8 {' '.join(map(str, cell))}\n")
        
        f.write(f"CELL_TYPES {len(cells)}\n")
        for _ in cells:
            f.write("12\n")
    
    print(f"VTK 文件已成功导出到 {output_file}")

if __name__ == "__main__":
    input_file = r"D:\360MoveData\Users\10758\Desktop\first_code\000035\q3_result\GridE_change.xlsx"
    output_file = r"D:\360MoveData\Users\10758\Desktop\first_code\000035\q3_result\output_lgr.vtk"
    print("加载数据")
    # 从Excel加载点数据
    geo_points, fault_points = load_points_from_excel(input_file)
    print("网格生成参数")
    # 调整网格生成参数
    target_cell_size = (10, 15)
    best_x_num, best_y_num, best_z_num, min_average_distance = optimize_grid(geo_points, fault_points, target_cell_size)
    print(f"最佳网格参数: x_num={best_x_num}, y_num={best_y_num}, z_num={best_z_num}, 最小平均距离: {min_average_distance}")

    # 生成网格节点
    grid_nodes = generate_grid_nodes(geo_points, best_x_num, best_y_num, best_z_num)

    # 使用最优参数重新构建六面体单元
    cells = construct_hexahedral_grid(best_x_num, best_y_num, best_z_num)

    # 导出为 VTK 文件
    export_to_vtk(grid_nodes, cells, output_file)